import matplotlib.pyplot as plt
import numpy as np

plt.rcParams.update({
    'font.family': 'serif',
    'font.serif': ['Times New Roman'],
    'font.size': 12,
    'axes.labelsize': 12,
    'axes.titlesize': 14,
    'xtick.labelsize': 10,
    'ytick.labelsize': 10,
    'legend.fontsize': 10,
    'figure.dpi': 300,
    'axes.grid.axis': 'y'  # 仅显示水平网格线
})

# 准备数据
models = ['resnet18/34', 'resnet34/50', 'resnet50/101', 'wrn16/28', 'resnet18/wrn16']
vanilla = [0.764, 0.802, 0.804, 0.861, 0.814]
cfd_clip = [0.954, 0.943, 0.965, 0.976, 0.957]

x = np.arange(len(models))  # 模型分类位置
width = 0.35  # 条形宽度

# 创建图表
fig, ax = plt.subplots(figsize=(8, 5))

# 绘制条形
rects1 = ax.bar(x - width/2, vanilla, width, label='Vanilla', color='#4C72B0', edgecolor='black')
rects2 = ax.bar(x + width/2, cfd_clip, width, label='CFD-CLIP', color='#55A868', edgecolor='black')

# 添加标签和标题
ax.set_ylabel('Similarity Statistics', fontweight='bold')
ax.set_title('Similarity Statistics of Features', fontweight='bold', pad=20)
ax.set_xticks(x)
ax.set_xticklabels(models, rotation=0, ha='center')  # 修改为水平居中
ax.legend(frameon=True, loc='upper left')

# 设置纵轴范围
ax.set_ylim(0, 1.0)  # 修改纵轴范围
ax.set_yticks(np.arange(0, 1.1, 0.1))  # 设置0.1间隔的刻度

# 添加数值标签
def autolabel(rects):
    for rect in rects:
        height = rect.get_height()
        ax.annotate(f'{height:.3f}',
                    xy=(rect.get_x() + rect.get_width() / 2, height),
                    xytext=(0, 3),
                    textcoords="offset points",
                    ha='center', 
                    va='bottom',
                    fontsize=8)

autolabel(rects1)
autolabel(rects2)

# 调整布局
plt.tight_layout()
plt.subplots_adjust(bottom=0.15)  # 增加底部边距
plt.savefig('similarity_stats.png', bbox_inches='tight')
plt.show()